#include "md501.h"
#include <arm_neon.h>
#include <cstring>

using namespace std;

Byte **StringProcessBatch(const string inputs[4], int n_bytes[4]) {
    Byte **paddedMessages = new Byte *[4];
    for (int idx = 0; idx < 4; ++idx) {
        if (inputs[idx].empty()) {
            paddedMessages[idx] = nullptr;
            n_bytes[idx] = 0;
            continue;
        }

        int length = inputs[idx].length();
        int bitLength = length * 8;
        int paddingBits = (bitLength + 64) % 512;
        paddingBits = (paddingBits < 448) ? (448 - paddingBits) : (512 - paddingBits + 448);

        int paddedLength = length + paddingBits / 8 + 8;
        paddedMessages[idx] = new Byte[paddedLength];

        memcpy(paddedMessages[idx], inputs[idx].c_str(), length);
        paddedMessages[idx][length] = 0x80;
        memset(paddedMessages[idx] + length + 1, 0, paddingBits / 8 - 1);

        // 写入原始长度（64位小端序）
        for (int i = 0; i < 8; ++i) {
            paddedMessages[idx][length + paddingBits / 8 + i] = ((uint64_t)bitLength >> (i * 8)) & 0xFF;
        }

        n_bytes[idx] = paddedLength;
    }
    return paddedMessages;
}

void MD5Hash(const string inputs[4], uint32_t states[4][4]) {
    // 初始化四路状态
    uint32x4_t abcd[4] = {
        vdupq_n_u32(0x67452301), // a
        vdupq_n_u32(0xefcdab89), // b
        vdupq_n_u32(0x98badcfe), // c
        vdupq_n_u32(0x10325476)  // d
    };

    int n_bytes[4];
    Byte **paddedMessages = StringProcessBatch(inputs, n_bytes);

    // 计算最大块数（所有输入统一处理）
    int max_blocks = 0;
    for (int idx = 0; idx < 4; ++idx) {
        if (paddedMessages[idx]) {
            max_blocks = max(max_blocks, n_bytes[idx] / 64);
        }
    }

    // 分块处理
    for (int block_idx = 0; block_idx < max_blocks; ++block_idx) {
        uint32x4_t x_simd[16];

        // 加载四路数据到SIMD寄存器
        for (int i = 0; i < 16; ++i) {
            uint32_t x[4] = {0};
            for (int idx = 0; idx < 4; ++idx) {
                if (paddedMessages[idx] && block_idx < n_bytes[idx] / 64) {
                    const Byte *block = paddedMessages[idx] + block_idx * 64;
                    x[idx] = block[4 * i] | (block[4 * i + 1] << 8) | 
                             (block[4 * i + 2] << 16) | (block[4 * i + 3] << 24);
                }
            }
            x_simd[i] = vld1q_u32(x); // 直接加载四路数据
        }

        // 保存初始状态
        uint32x4_t a = abcd[0], b = abcd[1], c = abcd[2], d = abcd[3];

        /* Round 1-4 运算（与原始代码相同，略） */
        // ... 使用 x_simd 进行FF/GG/HH/II_SIMD操作 ...

        // 更新状态
        abcd[0] = vaddq_u32(abcd[0], a);
        abcd[1] = vaddq_u32(abcd[1], b);
        abcd[2] = vaddq_u32(abcd[2], c);
        abcd[3] = vaddq_u32(abcd[3], d);
    }

    // 存储结果并调整字节序
    for (int idx = 0; idx < 4; ++idx) {
        uint32_t state[4];
        vst1q_u32(state, abcd[idx]);
        for (int i = 0; i < 4; ++i) {
            states[idx][i] = __builtin_bswap32(state[i]); // 使用编译器内置函数加速字节序转换
        }
    }

    // 释放内存
    for (int idx = 0; idx < 4; ++idx) {
        delete[] paddedMessages[idx];
    }
    delete[] paddedMessages;
}